{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Advanced constraints in static inverse free-boundary equilibrium calculations\n",
    "\n",
    "---\n",
    "\n",
    "In this example notebook, we demonstrate some of the more advanced types of constraints and techniques that can be employed in the inverse solver. \n",
    "\n",
    "#### Instatiate the objects\n",
    "\n",
    "As before, we start by instantiating the machine, equilibrium, profiles, and solver objects. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build machine\n",
    "from freegsnke import build_machine\n",
    "tokamak = build_machine.tokamak(\n",
    "    active_coils_path=\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n",
    "    passive_coils_path=\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n",
    "    limiter_path=\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n",
    "    wall_path=\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n",
    ")\n",
    "\n",
    "from freegsnke import equilibrium_update\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,      # provide tokamak object\n",
    "    Rmin=0.1, Rmax=2.0,   # radial range\n",
    "    Zmin=-2.2, Zmax=2.2,  # vertical range\n",
    "    nx=65,                # number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    ny=129,               # number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    # psi=plasma_psi\n",
    ")\n",
    "\n",
    "# initialise the profiles\n",
    "from freegsnke.jtor_update import ConstrainPaxisIp\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,        # equilibrium object\n",
    "    paxis=8e3,    # profile object\n",
    "    Ip=6e5,       # plasma current\n",
    "    fvac=0.5,     # fvac = rB_{tor}\n",
    "    alpha_m=1.8,  # profile function parameter\n",
    "    alpha_n=1.2   # profile function parameter\n",
    ")\n",
    "\n",
    "from freegsnke import GSstaticsolver\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# as before, let's fix the Solenoid current\n",
    "eq.tokamak.set_coil_current('Solenoid', 5000)\n",
    "eq.tokamak['Solenoid'].control = False  # ensures the current in the Solenoid is fixed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Normalised) poloidal flux constraints at specific locations\n",
    "\n",
    "In addition to the `null points`, `isoflux set`, and `coil_current_limits` constraints introduced in the previous notebook, additional methods are available for more advanced control of the inverse problem.\n",
    "\n",
    "We can also constrain:\n",
    "\n",
    "- **Flux values (`psi_vals`)**  \n",
    "  If the flux is known at a location $(R_j, Z_j)$, we can impose\n",
    "  $$\n",
    "  \\psi(R_j, Z_j) = \\psi_j^{\\text{target}}.\n",
    "  $$\n",
    "  More generally, one could prescribe $\\psi(R,Z)$ over a region (e.g. a full flux map). Flux values are typically difficult to estimate a priori to simulation and so the following constraint is often more useful. \n",
    "\n",
    "- **Normalised flux values (`psi_norm_limits`)**  \n",
    "  At a location $(R_j, Z_j)$, we can impose upper (or lower) bound constraints on the normalised flux \n",
    "  $$\n",
    "  \\hat{\\psi}(R,Z) = \\frac{\\psi(R,Z) - \\psi_{axis}}{\\psi_{boundary} - \\psi_{axis}},\n",
    "  $$ \n",
    "  such that \n",
    "  $$\n",
    "  \\hat{\\psi}(R_j, Z_j) \\leq \\hat{\\psi}_j^{\\text{target}}.\n",
    "  $$\n",
    "  or\n",
    "  $$\n",
    "  \\hat{\\psi}(R_j, Z_j) \\geq \\hat{\\psi}_j^{\\text{target}}.\n",
    "  $$\n",
    "  Note that $\\psi_{axis}$ and $\\psi_{boundary}$ are the values of the flux on the magnetic axis and plasma boundary, respectively.\n",
    "  As we'll see, this can be useful for setting explicit constraints on flux behaviour near the wall if there are certain safety limits to adhere to.\n",
    "\n",
    "Once again, we stress that users should be mindful of specifying too many constraints that may conflict with one another during an inverse solve. For example, a given normalised $\\psi$ constraint may be impossible/difficult to achieve given other isoflux constraints or coil current limits. More generally, specifying too many constraints could make the problem ill-posed and reduce the overall quality of the final solution. \n",
    "\n",
    "Recall that if you find that the solver is violating the limit (inequality) constraints (coil limits and normalised $\\psi$), you should look to reduce the number of constraints or try modifying the `mu_*` parameters. These parameters control how much a violation of the constraint is penalised in the solver, and increasing this penalty (by increasing `mu_*`) will increase the likelihood the constraint is adequately satisfied. Again, it may also be that the constraint is impossible to satisfy given your other constraints.\n",
    "\n",
    "In the following, we will show how to use the normalised flux constraints to control the normalised flux on the divertor nose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "from freegsnke.inverse import Inverse_optimizer\n",
    "\n",
    "Rx = 0.6      # X-point radius\n",
    "Zx = 1.1      # X-point height\n",
    "Rout = 1.4    # outboard midplane radius\n",
    "Rin = 0.34    # inboard midplane radius\n",
    "\n",
    "# set desired null_points locations (this can include X-point and O-point locations)\n",
    "null_points = [[Rx, Rx], [Zx, -Zx]]\n",
    "\n",
    "# Set desired isoflux constraints with format \n",
    "# isoflux_set = [isoflux_0, isoflux_1 ... ] \n",
    "# with each isoflux_i = [R_coords, Z_coords, weights]\n",
    "isoflux_set = np.array([\n",
    "    [\n",
    "        [Rx, Rx, Rin, Rout], \n",
    "        [Zx, -Zx, 0., 0.],\n",
    "    ]\n",
    "])\n",
    "\n",
    "# set the coil current limits (upper and lower)\n",
    "# coil ordering in this case: PX,  D1,  D2,  D3,  Dp,  D5,  D6,  D7,  P4,  P5,  P6\n",
    "coil_current_limits = [\n",
    "    [5e3, 9e3, 9e3, 7e3, 7e3, 5e3, 4e3, 5e3, 0.0, 0.0, None],\n",
    "    [-5e3, -9e3, -9e3, -7e3, -7e3, -5e3, -4e3, -5e3, -10e3, -10e3, None]\n",
    "]\n",
    "\n",
    "# normalised psi constraints set with format:\n",
    "# [R, Z, psiN, 1] --> ψ̂(R,Z) ≥ psiN  (the +1 or -1 defines ≥ or ≤)\n",
    "# [R, Z, psiN, -1] --> ψ̂(R,Z) ≤ psiN\n",
    "psi_norm_limits = [\n",
    "    [0.82, -1.55, 1.2, 1],\n",
    "    [0.82, -1.55, 1.3, -1],\n",
    "]\n",
    "\n",
    "# instantiate the freegsnke constrain object\n",
    "constrain = Inverse_optimizer(\n",
    "    null_points=null_points,\n",
    "    isoflux_set=isoflux_set,\n",
    "    coil_current_limits=coil_current_limits,\n",
    "    psi_norm_limits=psi_norm_limits,\n",
    "    mu_psi_norm = 1e7, # penalise how much the normalised psi constraint is violated (default 1e6)\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We solve, as in the previous example, by passing the equilibrium, profiles, and constraints to the static solver. Again, we apply regularisation to encourage a solution with low coil currents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GSStaticSolver.solve(eq=eq, \n",
    "                     profiles=profiles, \n",
    "                     constrain=constrain, \n",
    "                     target_relative_tolerance=1e-6,\n",
    "                     target_relative_psit_update=1e-3,\n",
    "                     verbose=True, # print output\n",
    "                     l2_reg=np.array([1e-12]*10+[1e-6]), \n",
    "                     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# plot the resulting equilbria \n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "constrain.plot(axis=ax1,show=True)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now check that the normalised $\\psi$ constraint has indeed been satisfied!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for psi_con in psi_norm_limits:\n",
    "    sign = \">=\" if psi_con[3] >= 0 else \"<=\"\n",
    "    print(f\"Psi norm at ({psi_con[0]}, {psi_con[1]}) = {eq.psiNRZ(psi_con[0], psi_con[1]):.3f} ({sign} {psi_con[2]})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Weighting constraints within the solver\n",
    "\n",
    "It is possible to specify **weights** on the different types of constraints within the inverse solver, enabling the solver to prioritise certain constraints over others.\n",
    "\n",
    "There are two ways of doing this:\n",
    "\n",
    "1. **Weighting different types of constraints relative to one another:**\n",
    "\n",
    "   Different classes of constraint can be weighted relative to each other (e.g. null points, isoflux, and psi constraints). This allows the solver to be instructed that one type of constraint is more important than others. For example, a user may want null point constraints to be strictly satisfied, but may be less concerned about the isoflux set being met with the same strictness.\n",
    "\n",
    "2. **Weighting individual isoflux constraints within a set:**\n",
    "\n",
    "   Individual isoflux constraints within a set can optionally be weighted relative to one another, instructing the solver to prioritise satisfying some subset over others.\n",
    "   \n",
    "   For example, consider a set of three isoflux constraints: one at the X-point, one at the outer midplane, and one at a desired strike point. The strike point constraint could be assigned a lower weight to allow the solver to find a solution with a strike point close to — but not necessarily exactly on — the target location."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Option 1: weighting different types of constraints\n",
    "\n",
    "Here, we'll show how to weight the null point constraints **more** than the isoflux constraints. This solver then prioritises fitting the null points exactly and the isoflux points less so. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# core constraints\n",
    "Rx = 0.55     # X-point radius\n",
    "Zx = 1.2      # X-point height\n",
    "Rout = 1.4    # outboard midplane radius\n",
    "Rin = 0.34    # inboard midplane radius\n",
    "\n",
    "# set desired null_points locations (this can include X-point and O-point locations)\n",
    "null_points = [[Rx, Rx], [Zx, -Zx]]\n",
    "\n",
    "# set desired isoflux constraints with format \n",
    "# isoflux_set = [isoflux_0, isoflux_1 ... ] \n",
    "# with each isoflux_i = [R_coords, Z_coords]\n",
    "isoflux_set = np.array([\n",
    "    [\n",
    "        [Rx, Rx, Rin, Rout, 0.75, 1.0, 0.75, 1.0], \n",
    "        [Zx, -Zx, 0.0, 0.0, -1.6, -2.0, 1.6, 2.0],\n",
    "    ]\n",
    "])\n",
    "\n",
    "# instantiate the freegsnke constrain object\n",
    "constrain = Inverse_optimizer(\n",
    "    null_points=null_points,\n",
    "    isoflux_set=isoflux_set,\n",
    "    coil_current_limits=coil_current_limits,\n",
    "    weight_isoflux=0.2, # <-- weight the isofluxes less than the null points (default 1.0)\n",
    "    weight_nulls=1.0,   # <-- weight the null points more than the null points (default 1.0)\n",
    "    # weight_psi=1.0    # <-- weight for psi values not used here\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# solve!\n",
    "GSStaticSolver.solve(eq=eq, \n",
    "                     profiles=profiles, \n",
    "                     constrain=constrain, \n",
    "                     target_relative_tolerance=1e-6,\n",
    "                     target_relative_psit_update=1e-3,\n",
    "                     verbose=True,\n",
    "                     l2_reg=np.array([1e-12]*10+[1e-6]), \n",
    "                     )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the null points constraints are satisfied pefectly while the isoflux constraints are \"less so\" (especially in the divertor region)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the resulting equilbria \n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "constrain.plot(axis=ax1,show=True)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Option 2: weighting different isoflux constraints\n",
    "\n",
    "Here, we'll show how to weight individual isoflux constraints differently from one another. For example, here we want to prioritise the core shape over the exact location of the strikepoint. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# core constraints\n",
    "Rx = 0.55     # X-point radius\n",
    "Zx = 1.2      # X-point height\n",
    "Rout = 1.4    # outboard midplane radius\n",
    "Rin = 0.34    # inboard midplane radius\n",
    "\n",
    "# set desired null_points locations (this can include X-point and O-point locations)\n",
    "null_points = [[Rx, Rx], [Zx, -Zx]]\n",
    "\n",
    "# set desired isoflux constraints with format \n",
    "# isoflux_set = [isoflux_0, isoflux_1 ... ] \n",
    "# with each isoflux_i = [R_coords, Z_coords, weights]\n",
    "isoflux_set = np.array([\n",
    "    [\n",
    "        [Rx, Rx, Rin, Rout, 1.0, 1.0, 0.75, 1.0, 0.75, 1.0], \n",
    "        [Zx, -Zx, 0.0, 0.0, -0.9, 0.9,  -1.6, -2.0, 1.6, 2.0],\n",
    "        [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.1, 1.0, 0.1],\n",
    "    ]\n",
    "])\n",
    "\n",
    "# instantiate the freegsnke constrain object\n",
    "constrain = Inverse_optimizer(\n",
    "    null_points=null_points,\n",
    "    isoflux_set=isoflux_set,\n",
    "    coil_current_limits=coil_current_limits,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# solve!\n",
    "GSStaticSolver.solve(eq=eq, \n",
    "                     profiles=profiles, \n",
    "                     constrain=constrain, \n",
    "                     target_relative_tolerance=1e-6,\n",
    "                     target_relative_psit_update=1e-3,\n",
    "                     verbose=True, # print output\n",
    "                     l2_reg=np.array([1e-12]*10+[1e-6]), \n",
    "                     )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the strikepoint isoflux constraint is no longer satisfied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the resulting equilbria \n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "constrain.plot(axis=ax1,show=True)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.10.17)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}